from typing import Any, Dict, List, Optional, Tuple
from typing_extensions import Self

from datetime import datetime
import numpy as np
import os
import pickle
from rdkit import RDLogger
import re
RDLogger.DisableLog("rdApp.*")
from scipy import spatial

from open_biomed.data.molecule import Molecule
from open_biomed.data.protein import Protein

POCKET_CONFIG = {
    'bounds': [26.52218454781356, 27.875010450039014, 28.827148928234973, 29.653040886224233, 30.439636480866106, 31.234028070512103, 32.014619850535716, 33.035326071776495, 34.57565983834558], 
    'bins': [
        ([10, 6, 9, 22, 11, 13, 12, 28, 8, 7, 16, 18, 5, 19, 17, 20, 24, 27, 14, 26, 21, 15, 33, 29, 3, 4, 30, 31, 23, 34, 36, 25, 32, 39, 35, 37, 38, 40, 43, 2, 41],[0.1196, 0.0839, 0.1152, 0.0065, 0.1131, 0.0645, 0.0807, 0.0086, 0.0807, 0.0626, 0.0151, 0.019, 0.0092, 0.024, 0.0205, 0.0164, 0.011, 0.0122, 0.03, 0.0141, 0.0105, 0.0164, 0.0033, 0.0061, 0.0031, 0.0085, 0.0052, 0.0057, 0.0146, 0.002, 0.0019, 0.0084, 0.0024, 0.0007, 0.0017, 0.0012, 0.0005, 0.0003, 0.0003, 0.0002, 0.0001]),
        ([26, 13, 28, 22, 20, 15, 11, 10, 25, 14, 9, 12, 17, 34, 24, 21, 19, 16, 23, 43, 32, 18, 51, 8, 30, 33, 29, 31, 27, 7, 38, 6, 37, 35, 36, 44, 40, 39, 5, 50, 41, 53], [0.0174982501749825, 0.0746925307469253, 0.0184981501849815, 0.0168983101689831, 0.0443955604439556, 0.043695630436956304, 0.12038796120387961, 0.08189181081891811, 0.017698230176982303, 0.0730926907309269, 0.04559544045595441, 0.10388961103889612, 0.037396260373962605, 0.0047995200479952005, 0.0170982901709829, 0.029697030296970302, 0.0357964203579642, 0.040695930406959305, 0.024697530246975304, 0.0007999200079992001, 0.013098690130986902, 0.033396660333966606, 0.00019998000199980003, 0.0202979702029797, 0.0112988701129887, 0.007799220077992201, 0.0115988401159884, 0.009799020097990201, 0.0187981201879812, 0.007799220077992201, 0.00039996000399960006, 0.0037996200379962005, 0.0011998800119988001, 0.005699430056994301, 0.001999800019998, 0.0008999100089991, 0.00029997000299970003, 0.0018998100189981002, 0.00019998000199980003, 9.999000099990002e-05, 9.999000099990002e-05, 9.999000099990002e-05]),
        ([16, 32, 24, 18, 23, 20, 21, 14, 27, 12, 11, 19, 30, 22, 29, 15, 9, 33, 25, 13, 28, 42, 26, 17, 10, 37, 31, 35, 44, 40, 38, 34, 8, 54, 36, 39, 6, 7, 45, 51, 46, 43, 5, 41], [ 0.0604120824164833, 0.010502100420084017, 0.032906581316263256, 0.05471094218843769, 0.05131026205241048, 0.07371474294858972, 0.0666133226645329, 0.061312262452490496, 0.03680736147229446, 0.04380876175235047, 0.03170634126825365, 0.057411482296459294, 0.017603520704140826, 0.042008401680336066, 0.02430486097219444, 0.04740948189637927, 0.010202040408081616, 0.004100820164032806, 0.03790758151630326, 0.03470694138827766, 0.027005401080216044, 0.00010002000400080016, 0.032906581316263256, 0.0614122824564913, 0.03700740148029606, 0.000600120024004801, 0.011702340468093618, 0.011302260452090419, 0.00020004000800160032, 0.0018003600720144029, 0.0014002800560112022, 0.003800760152030406, 0.003800760152030406, 0.00010002000400080016, 0.003200640128025605, 0.0007001400280056011, 0.0011002200440088017, 0.0014002800560112022, 0.00040008001600320064, 0.00020004000800160032, 0.00010002000400080016, 0.00010002000400080016, 0.00010002000400080016, 0.00010002000400080016]),
        ([24, 18, 10, 28, 25, 20, 17, 19, 22, 21, 26, 23, 13, 11, 15, 27, 29, 31, 16, 30, 14, 12, 35, 37, 9, 39, 34, 33, 40, 36, 38, 32, 44, 51, 42, 43, 8, 41, 48, 47, 57, 50, 6], [ 0.04357821089455272, 0.04717641179410295, 0.030584707646176913, 0.031784107946026985, 0.048875562218890556, 0.088055972013993, 0.044077961019490255, 0.06426786606696652, 0.051274362818590706, 0.06966516741629185, 0.043778110944527736, 0.05907046476761619, 0.021289355322338832, 0.009695152423788106, 0.041379310344827586, 0.04737631184407796, 0.03158420789605197, 0.02258870564717641, 0.06436781609195402, 0.024387806096951525, 0.0336831584207896, 0.0223888055972014, 0.012393803098450775, 0.0025987006496751624, 0.0023988005997001498, 0.0005997001499250374, 0.006596701649175412, 0.009495252373813094, 0.0026986506746626685, 0.0063968015992004, 0.002098950524737631, 0.009895052473763119, 0.0005997001499250374, 9.995002498750625e-05, 0.0010994502748625688, 0.0004997501249375312, 0.0002998500749625187, 0.0005997001499250374, 0.0001999000499750125, 0.0001999000499750125, 9.995002498750625e-05, 9.995002498750625e-05, 9.995002498750625e-05]),
        ([17, 14, 25, 24, 32, 29, 18, 20, 28, 19, 41, 23, 31, 15, 35, 27, 34, 44, 30, 13, 9, 16, 26, 21, 38, 33, 22, 36, 10, 37, 12, 57, 11, 39, 50, 47, 40, 42, 55, 67, 48, 43, 45], [ 0.04202101050525263, 0.024712356178089045, 0.0608304152076038, 0.04692346173086543, 0.0208104052026013, 0.037818909454727365, 0.04202101050525263, 0.10215107553776888, 0.045822911455727866, 0.05542771385692846, 0.0030015007503751876, 0.06303151575787894, 0.03071535767883942, 0.03251625812906453, 0.011705852926463232, 0.05882941470735368, 0.011005502751375688, 0.0008004002001000501, 0.0304152076038019, 0.009804902451225613, 0.0008004002001000501, 0.040720360180090044, 0.05042521260630315, 0.08584292146073036, 0.0026013006503251626, 0.01310655327663832, 0.04392196098049025, 0.0073036518259129565, 0.0026013006503251626, 0.003901950975487744, 0.008204102051025513, 0.00010005002501250626, 0.0034017008504252126, 0.0010005002501250625, 0.0005002501250625312, 0.00030015007503751874, 0.003501750875437719, 0.0006003001500750375, 0.00010005002501250626, 0.00010005002501250626, 0.00040020010005002504, 0.00010005002501250626, 0.00010005002501250626]),
        ([26, 38, 17, 24, 32, 25, 34, 28, 31, 22, 21, 29, 30, 20, 23, 19, 42, 27, 65, 35, 18, 16, 41, 14, 33, 37, 43, 13, 10, 15, 36, 40, 48, 11, 39, 12, 44, 9, 49, 66, 45, 67, 69, 51, 57, 55, 50], [0.0512, 0.0046, 0.0297, 0.0525, 0.0356, 0.0624, 0.0206, 0.0541, 0.0379, 0.0739, 0.0829, 0.0397, 0.0424, 0.0801, 0.0643, 0.0421, 0.0011, 0.0656, 0.001, 0.018, 0.0326, 0.0236, 0.0024, 0.0109, 0.0238, 0.0047, 0.0009, 0.0034, 0.001, 0.0134, 0.0092, 0.0029, 0.001, 0.0017, 0.0026, 0.002, 0.001, 0.0012, 0.0003, 0.0003, 0.0002, 0.0002, 0.0001, 0.0001, 0.0006, 0.0001, 0.0001]),
        ([26, 28, 29, 22, 19, 23, 24, 30, 21, 15, 33, 36, 35, 25, 32, 34, 31, 27, 20, 14, 17, 16, 40, 58, 13, 41, 39, 38, 18, 48, 37, 44, 43, 11, 42, 50, 57, 45, 54, 46, 52, 65, 49, 10, 69, 12, 66, 47], [0.0613, 0.0816, 0.0496, 0.0506, 0.0245, 0.0567, 0.0565, 0.0562, 0.0537, 0.0079, 0.035, 0.0186, 0.0218, 0.0664, 0.0488, 0.0215, 0.059, 0.0912, 0.0416, 0.0021, 0.0155, 0.0124, 0.0045, 0.0004, 0.0014, 0.0042, 0.0025, 0.0089, 0.0225, 0.0012, 0.0088, 0.001, 0.0027, 0.0001, 0.002, 0.0004, 0.0014, 0.001, 0.0003, 0.001, 0.0002, 0.0012, 0.0004, 0.0004, 0.0002, 0.0005, 0.0002, 0.0001]),
        ([34, 32, 27, 22, 20, 35, 33, 23, 28, 29, 26, 24, 21, 31, 25, 16, 40, 17, 19, 37, 44, 30, 36, 41, 43, 38, 10, 42, 15, 18, 57, 39, 48, 51, 66, 52, 54, 46, 53, 45, 50, 12, 49, 67, 13, 14, 65, 11, 47, 58, 69], [0.0293, 0.0612, 0.0847, 0.0332, 0.0387, 0.0219, 0.0424, 0.0557, 0.075, 0.0449, 0.0627, 0.0486, 0.0342, 0.0693, 0.0656, 0.0111, 0.0144, 0.0124, 0.0236, 0.0122, 0.003, 0.0585, 0.0234, 0.0093, 0.0048, 0.0155, 0.0002, 0.0087, 0.0041, 0.0101, 0.0007, 0.007, 0.0007, 0.0004, 0.0003, 0.0003, 0.0007, 0.0029, 0.0001, 0.0023, 0.0004, 0.0004, 0.0009, 0.0001, 0.0004, 0.0025, 0.0003, 0.0004, 0.0003, 0.0001, 0.0001]),
        ([32, 37, 26, 46, 28, 33, 27, 31, 30, 29, 34, 38, 21, 25, 24, 22, 36, 35, 41, 17, 23, 40, 39, 48, 42, 63, 45, 20, 15, 43, 54, 50, 12, 19, 44, 13, 58, 47, 51, 67, 57, 18, 56, 16, 49, 61, 55, 68, 52, 62, 11, 14, 53, 10, 59], [0.0626, 0.023, 0.0854, 0.0026, 0.0624, 0.0632, 0.0835, 0.0733, 0.0678, 0.0509, 0.0506, 0.0197, 0.022, 0.0486, 0.04, 0.024, 0.0332, 0.0416, 0.0097, 0.0034, 0.0368, 0.0127, 0.0126, 0.0031, 0.012, 0.0002, 0.0038, 0.0105, 0.0026, 0.0062, 0.0052, 0.0005, 0.0002, 0.0087, 0.0045, 0.0002, 0.0008, 0.0021, 0.0009, 0.0005, 0.001, 0.0027, 0.0003, 0.0012, 0.001, 0.0003, 0.0002, 0.0001, 0.0006, 0.0001, 0.0001, 0.0004, 0.0002, 0.0001, 0.0001]),
        ([35, 44, 49, 26, 39, 32, 23, 31, 29, 27, 45, 33, 38, 28, 40, 42, 36, 24, 67, 30, 22, 43, 34, 50, 19, 37, 51, 41, 17, 48, 25, 21, 66, 13, 16, 46, 20, 55, 52, 47, 15, 18, 54, 63, 65, 56, 57, 60, 58, 61, 70, 59, 64, 53, 14, 12, 69, 86], [0.0521, 0.0663, 0.0027, 0.0381, 0.0427, 0.0696, 0.0094, 0.0604, 0.0383, 0.0423, 0.0087, 0.0815, 0.0379, 0.0491, 0.0163, 0.0175, 0.0474, 0.0367, 0.0022, 0.0413, 0.0095, 0.022, 0.0513, 0.0016, 0.0026, 0.0399, 0.0017, 0.0285, 0.0011, 0.0101, 0.0218, 0.0077, 0.0019, 0.0002, 0.0014, 0.0096, 0.007, 0.0025, 0.0025, 0.0066, 0.0005, 0.0022, 0.0017, 0.0001, 0.001, 0.0009, 0.0015, 0.0001, 0.0002, 0.0002, 0.0002, 0.0003, 0.0002, 0.0004, 0.0001, 0.0001, 0.0002, 0.0001])
    ]
}

class Pocket:
    def __init__(self, name: Optional[str] = None) -> None:
        self.atoms = None
        self.conformer = None
        self.orig_indices = None         # Which part of the protein it belongs to
        self.orig_protein = None
        self.name = name

    @classmethod
    def from_protein_subregion(cls, protein: Protein, center: List[float], radius: float=10.0) -> Self:
        raise NotImplementedError

    @classmethod
    def from_protein_ref_ligand(cls, protein: Protein, ligand: Molecule, radius: float=10.0) -> Self:
        sel_idx = []
        for pos in ligand.conformer:
            for i, residue in enumerate(protein.residues):
                dist = np.linalg.norm(residue.center_of_mass - pos, ord=2, axis=-1)
                if dist < radius and not i in sel_idx:
                    sel_idx.append(i)
        return Pocket.from_protein_subseq(protein, sel_idx)

    @classmethod
    def from_protein_subseq(cls, protein: Protein, indices: List[int]) -> Self:
        pocket = cls()
        pocket.atoms, pocket.conformer, pocket.orig_indices = [], [], []
        for i in indices:
            for atom_i in protein.residues[i].atoms:
                atom = protein.all_atom[atom_i]
                pocket.atoms.append(atom)
                pocket.conformer.append(atom['pos'])
            pocket.orig_indices.append(i)
        pocket.conformer = np.array(pocket.conformer)
        pocket.orig_protein = protein
        return pocket
    
    @classmethod
    def from_binary_file(cls, file: str) -> Self:
        return pickle.load(open(file, "rb"))

    @classmethod
    def from_pdb_file(cls, pdb_file: str, removeHs: bool=True) -> Self:
        protein = Protein.from_pdb_file(pdb_file, removeHs=removeHs)
        pocket = cls.from_protein_subseq(protein, list(range(len(protein.residues))))
        pocket.orig_protein = protein
        return pocket
    
    def save_pdb(self, file: Optional[str]=None, overwrite: bool=False) -> str:
        if file is None:
            self._add_name()
            file = f"./tmp/{self.name}.pdb"

        return self.orig_protein.save_pdb(file, overwrite, self.orig_indices)

    def _add_name(self) -> None:
        if self.name is None:
            self.name = "pocket_" + re.sub(r"[-:.]", "_", datetime.now().isoformat(sep="_", timespec="milliseconds"))

    def save_binary(self, file: Optional[str]=None, overwrite: bool=False) -> str:
        if file is None:
            self._add_name()
            file = f"./tmp/{self.name}.pkl"

        if not os.path.exists(file) or overwrite:
            pickle.dump(self, open(file, "wb"))
        return file
    
    def get_num_atoms(self) -> int:
        return len(self.atoms)

    def __str__(self) -> str:
        if self.orig_protein is not None:
            return "A pocket from protein " + str(self.orig_protein) + " at positions " + ",".join([str(idx+1) for idx in self.orig_indices])
        else:
            return "A pocket"

def estimate_ligand_atom_num(pocket: Pocket) -> int:
    if hasattr(pocket, "estimated_num_atoms"):
        return pocket.estimated_num_atoms
    dist = spatial.distance.pdist(pocket.conformer, metric='euclidean')
    dist = np.sort(dist)[::-1]
    space_size = np.median(dist[:10])
    bounds = POCKET_CONFIG['bounds']
    for bin_idx in range(len(bounds)):
        if bounds[bin_idx] > space_size:
            break
    num_atoms, probs = POCKET_CONFIG['bins'][bin_idx]
    return np.random.choice(num_atoms, p=probs)